# -*- coding: utf-8 -*-
"""
Created on Mon Sep 28 16:57:26 2020

@author: Robin
"""

import numpy as np
import pandas as pd
from .performance import factor_information_coefficient
from .utils import (cal_factor_data,merge_to_multiindex)
from sklearn.linear_model import LinearRegression
from .utils import compute_forward_returns

def get_combined_factor_df(factors,stock,model_type = 'ic_ir'):
    """
    传入因子数据、股票数据、
    计算日期、回溯时长
    """
    weights = {}
    # 计算因子权重
    for ind_code in factors.columns:
        # 生成对齐数据
        factor_data = cal_factor_data(stock,factors[ind_code])
        # 计算ic-ir作为权重
        ictmp = factor_information_coefficient(factor_data[{'1D','factor'}],group_adjust=False,by_group=False)
        ic_ir = (ictmp.mean() / ictmp.std()).values[0]   # 越大越好
        weights[ind_code] = ic_ir
           
    # 因子合成
    combined_factor_df = sum(factors[ind_code] * weights[ind_code] for ind_code in factors.columns)
    
    return combined_factor_df

def get_top_stocks(df,stocknum):
    return df.sort_values(ascending = False)[:stocknum].index.values   
        
def get_stock_list_by_score(factor,groupby,datestr,stock_num):
    '''
    打分法，根据单个日期进行选股
    '''
    factor_name = factor.name            
    date_dt = pd.to_datetime(datestr)           
    dateindex = factor.index.levels[0]        
    select_date = dateindex[dateindex<=date_dt][-1]
    
    if (groupby is not None):
        # 分组
        factor = merge_to_multiindex(factor,groupby)
        select_data = factor.loc[select_date]     #series可以不用loc，但是dataframe必须用loc
        stock_df = select_data.groupby('industry')[factor_name].apply(get_top_stocks,stock_num)
        # 返回结果
        stock_choosed = stock_df.to_dict()
        stock_choosed_type = 2      
    else:   
        # 不分组
        select_data = factor.loc[select_date]     #series可以不用loc，但是dataframe必须用loc        
        # 排序筛选
        stock_list = select_data.sort_values(ascending = False)[:stock_num].index.values  
        # 返回结果
        stock_choosed = list(stock_list)   
        stock_choosed_type = 1
    return stock_choosed,stock_choosed_type

def get_stock_list_by_score_by_dates(factors,stock,groupby,datelist,stock_num,lookbackPeriod,combine_type):
    '''
    打分法，根据多个日期进行选股
    '''
    stock_choosed = {}  
    if (groupby is not None):
        # 分组        
        for datetmp in datelist:
            # 1.根据回望期以及日期序列合并相关因子，得到合并后的因子
            datetmp1 = pd.to_datetime(datetmp) # datetmp[:4]+'-'+datetmp[4:6]+'-'+datetmp[6:]
            # 提取相应日期的数据
            date_series_tmp = stock.index[stock.index<=datetmp1][-lookbackPeriod:]
            factorstmp = factors.loc[date_series_tmp]
            stocktmp = stock.loc[date_series_tmp]                
            # 因子合成
            combined_factor_df = get_combined_factor_df(factorstmp,stocktmp,combine_type)
            combined_factor_df.name = 'combine_factor'
            factor_name = 'combine_factor'
            select_date = date_series_tmp[-1] 

            # 分组
            factor = merge_to_multiindex(combined_factor_df,groupby)
            select_data = factor.loc[select_date]     #series可以不用loc，但是dataframe必须用loc
            # 排序筛选
            stock_df = select_data.groupby('industry')[factor_name].apply(get_top_stocks,stock_num)
            # 选股结果
            stock_list = stock_df.to_dict()   
            stock_choosed[datetmp] = stock_list    
            
        stock_choosed_type = 4
    else:
        # 不分组
        for datetmp in datelist:
            '''
            # 按照factor里面的行业信息进行选股，设置一个是否需要使用行业信息来选股的参数
            '''
            # 1.根据回望期以及日期序列合并相关因子，得到合并后的因子
            datetmp1 = pd.to_datetime(datetmp) # datetmp[:4]+'-'+datetmp[4:6]+'-'+datetmp[6:]
            # 提取相应日期的数据
            date_series_tmp = stock.index[stock.index<=datetmp1][-lookbackPeriod:]
            factorstmp = factors.loc[date_series_tmp]
            stocktmp = stock.loc[date_series_tmp]
            
            # 因子合成
            combined_factor_df = get_combined_factor_df(factorstmp,stocktmp,combine_type)
            combined_factor_df.name = 'combine_factor'
            factor_name = 'combine_factor'
            
            select_date = date_series_tmp[-1] 
            # 不分组
            select_data = combined_factor_df.loc[select_date]     #series可以不用loc，但是dataframe必须用loc        
            # 排序筛选
            stock_list = select_data.sort_values(ascending = False)[:stock_num].index.values  
            # 返回选股结果
            stock_list = list(stock_list) 
            stock_choosed[datetmp] = stock_list     
                   
        stock_choosed_type = 3
        
    return stock_choosed,stock_choosed_type

def get_stock_list_by_regression(factors,stock,groupby,datestr,stock_num):
    '''
    这里的factors格式同am.factors，即双索引（date,asset）的DataFrame
    
    date          asset     因子1    因子2   因子3
    2020-01-03    股票1      0.43    -0.23   1.2
                  股票2      0.33    -0.10   0.98
                  ...        ...     ...     ...
    2020-01-10    股票1      1.8     0.52    1.01
                  股票2      1.33    -0.90   0.38
                  ...        ...     ...     ...


    '''
    
    if (groupby is not None):
        # 分组
        stock_choosed_type = 2      
    else:     
        stock_choosed_type = 1
    
    select_ind = list(factors.columns)
    forward_returns = compute_forward_returns(factors[select_ind[0]],stock)
    
    date_dt = pd.to_datetime(datestr)
    dateindex = factors.index.levels[0]        
    select_date = dateindex[dateindex<=date_dt][-1]
    
    #用day1的收益率数据作为因变量进行回归
    ret_series=forward_returns.loc[select_date]['1D']        
    #用day0的因子数据作为自变量进行回归
    indi_date=dateindex[dateindex<=date_dt][-2]        
    indi_df=factors.loc[indi_date]
    
    #多元线性回归
    Y=ret_series.fillna(0)
    X=indi_df.fillna(0)
    
    lr = LinearRegression()
    lr.fit(X,Y)
    
    #用day1的因子值预测day2的收益率数据
    X_use=factors.loc[select_date]
    Y_pred = lr.predict(X_use)
    
    predict_ret=pd.Series(Y_pred,index=stock.columns)
    stock_list = predict_ret.sort_values(ascending = False)[:stock_num].index.values                 
            
    # 返回结果
    stock_choosed = list(stock_list) 

    return stock_choosed,stock_choosed_type

def get_stock_list_by_regress_by_dates(factors,stock,groupby,datelist,stock_num,lookbackPeriod):
    stock_choosed = {}
    
    if (groupby is not None):
        # 分组
        stock_choosed_type = 2   
    else:     
        stock_choosed_type = 1
    
    select_ind = list(factors.columns)
    forward_returns = compute_forward_returns(factors[select_ind[0]],stock)
    dateindex = factors.index.levels[0]

    for datetmp in datelist:
        select_date = dateindex[dateindex<=datetmp][-1]
    
        #用day1的收益率数据作为因变量进行回归
        ret_series=forward_returns.loc[select_date]['1D']        
        #用day0的因子数据作为自变量进行回归
        indi_date=dateindex[dateindex<=datetmp][-2]        
        indi_df=factors.loc[indi_date]
    
        #多元线性回归
        Y=ret_series.fillna(0)
        X=indi_df.fillna(0)
    
        lr = LinearRegression()
        lr.fit(X,Y)
        
        #用day1的因子值预测day2的收益率数据
        X_use=factors.loc[select_date]
        Y_pred = lr.predict(X_use)
        
        predict_ret=pd.Series(Y_pred,index=stock.columns)
        stock_list = predict_ret.sort_values(ascending = False)[:stock_num].index.values                 
                    
        # 返回结果
        stock_choosed[datetmp] = list(stock_list) 
    
    return stock_choosed,stock_choosed_type
